package com.zhxq.ai.repository;

import com.zhxq.ai.domain.ChatHistory;
import com.zhxq.ai.domain.ChatMessage;
import com.zhxq.ai.domain.vo.ChatHistoryItemVo;
import com.zhxq.ai.domain.vo.ChatHistoryVo;
import com.zhxq.common.utils.SecurityUtils;
import com.zhxq.common.utils.UserThreadLocal;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.stereotype.Component;

import java.time.LocalDate;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Component("mongoChatHistoryRepository")
public class MongoChatHistoryRepository implements ChatHistoryRepository {

    @Autowired
    private MongoTemplate mongoTemplate;

    @Override
    public void save(String chatType, String chatId) {
        Query query = new Query();
        query.addCriteria(Criteria.where("chatType").is(chatType));
        Long userId = UserThreadLocal.getUserId();
        if(userId == null) {
            userId = SecurityUtils.getUserId();
        }
        query.addCriteria(Criteria.where("userId").is(userId));
        ChatHistory chatHistory = mongoTemplate.findOne(query, ChatHistory.class);
        if(chatHistory == null) {
            chatHistory = new ChatHistory();
            chatHistory.setUserId(userId);
            chatHistory.setChatType(chatType);
            List<String> chatIds = new ArrayList<>();
            chatIds.add(chatId);
            chatHistory.setChatIds(chatIds);
            mongoTemplate.insert(chatHistory);
        } else {
            if(chatHistory.getChatIds().contains(chatId)) {
                return;
            }
            chatHistory.getChatIds().add(chatId);
            mongoTemplate.save(chatHistory);
        }
    }

    @Override
    public List<String> getChatIds(String chatType) {
        return List.of();
    }

    @Override
    public List<ChatHistoryVo> getChatHistory(String chatType) {
        Query query = new Query();
        query.addCriteria(Criteria.where("chatType").is(chatType));
        Long userId = UserThreadLocal.getUserId();
        if(userId == null) {
            userId = SecurityUtils.getUserId();
        }
        query.addCriteria(Criteria.where("userId").is(userId));
        // 获取该用户对应的会话历史记录
        ChatHistory chatHistory = mongoTemplate.findOne(query, ChatHistory.class);
        if(chatHistory == null) {
            return List.of();
        }
        // 获取该用户对应的所有会话ID
        List<String> chatIds = chatHistory.getChatIds();
        List<ChatMessage> chatSessions = chatIds.stream().map(chatId -> {
            Query chatQuery = new Query();
            chatQuery.addCriteria(Criteria.where("chatId").is(chatId));
            return mongoTemplate.findOne(chatQuery, ChatMessage.class);
        }).toList();
        // 将这些会话根据日期分组
        Map<LocalDate, List<ChatMessage>> collect = chatSessions.stream().collect(Collectors.groupingBy(ChatMessage -> {
            return LocalDate.ofInstant(ChatMessage.getUpdateTime().toInstant(), ZoneId.systemDefault());
        }));
        List<ChatHistoryVo> chatHistoryVos = new ArrayList<>();
        // 遍历每个日期,将数据转化为ChatHistoryVo
        collect.forEach((time, list) -> {
            ChatHistoryVo chatHistoryVo = new ChatHistoryVo();
            // 设置日期
            chatHistoryVo.setUpdateTime(time);
            // 遍历每个会话，将每个chatMessage转化为ChatHistoryItemVo
            List<ChatHistoryItemVo> chatHistoryItemVos = list.stream().map(chatMessage -> {
                ChatHistoryItemVo chatHistoryItemVo = new ChatHistoryItemVo();
                chatHistoryItemVo.setChatId(chatMessage.getChatId());
                if (chatMessage.getMessageList() != null && !chatMessage.getMessageList().isEmpty()) {
                    chatHistoryItemVo.setMessage(chatMessage.getMessageList().get(0).getContent());
                    int lastIndex = chatMessage.getMessageList().size() - 1;
                    chatHistoryItemVo.setChatTime(chatMessage.getMessageList().get(lastIndex).getChatTime());
                }
                return chatHistoryItemVo;
            }).collect(Collectors.toList());
            chatHistoryVo.setChatHistoryItemVos(chatHistoryItemVos);
            chatHistoryVos.add(chatHistoryVo);
        });
        return chatHistoryVos;
    }
}
